import logging
import time
from typing import Any, Dict
import torch
import numpy as np

def setup_logger(name: str, log_level=logging.INFO) -> logging.Logger:
    """Setup logger for PocketFlow.
    
    Args:
        name: Logger name.
        log_level: Logging level.
        
    Returns:
        Configured logger instance.
    """
    logger = logging.getLogger(name)
    logger.setLevel(log_level)
    
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    
    ch = logging.StreamHandler()
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    
    return logger

def load_model(model_path: str, device='cpu') -> Any:
    """Load PyTorch model from file.
    
    Args:
        model_path: Path to model file.
        device: Device to load model to.
        
    Returns:
        Loaded model.
    """
    return torch.load(model_path, map_location=device)

def save_model(model: Any, model_path: str):
    """Save PyTorch model to file.
    
    Args:
        model: Model to save.
        model_path: Path to save model to.
    """
    torch.save(model, model_path)

class ProgressTracker:
    """Track and display progress of compression/training."""
    
    def __init__(self, total: int, name: str = 'Progress'):
        """Initialize progress tracker.
        
        Args:
            total: Total number of steps.
            name: Name of the progress tracker.
        """
        self.total = total
        self.current = 0
        self.name = name
        self.start_time = time.time()
        
    def update(self, n: int = 1):
        """Update progress.
        
        Args:
            n: Number of steps completed.
        """
        self.current += n
        elapsed = time.time() - self.start_time
        remaining = (self.total - self.current) * elapsed / max(1, self.current)
        
        print(f"\r{self.name}: {self.current}/{self.total} "
              f"[{elapsed:.1f}s<{remaining:.1f}s]", end="")
        
    def close(self):
        """Finish progress tracking."""
        elapsed = time.time() - self.start_time
        print(f"\r{self.name}: {self.current}/{self.total} "
              f"[{elapsed:.1f}s]")